1499F - Diameter Cuts - CodeForces Solution

#include <bits/stdc++.h>

#define ff first

#define ss second

#define endl '\n'

using namespace std;

const long long INF = (long long) 1e18;

const int mod = (int) 998'244'353;

const int MAXN = (int) 5e3+5;



typedef long long ll;

typedef unsigned long long ull;

typedef pair<int,int> pii;

typedef pair<ll,ll> pll;

int n, k;



vector<int> adj[MAXN];

int dp[MAXN][MAXN]; // dp[i][j] i. node'da en fazla j derinlikli kaç durum var

int dep[MAXN];

void dfs(int v, int par){

    dep[v] = 0;

    

    for(int i = 0; i < adj[v].size(); i++){

        if(adj[v][i] == par){

            swap(adj[v][i], adj[v].back());

            adj[v].pop_back();

            break;

        }

    }



    for(int j: adj[v]){

        if(j == par)

            continue;

        dfs(j, v);

        dep[v] = max(dep[v], dep[j] + 1);

    }



    dep[v] = min(dep[v], k);

    dp[v][0] = 1;



    if(adj[v].empty())

        return;



    sort(adj[v].begin(), adj[v].end(), [](int a, int b){

        return dep[a] > dep[b];

    });

    int top = 0;

    for(int i = 0; i <= dep[adj[v][0]]; i++){

        dp[v][i + 1] = dp[adj[v][0]][i];

        top = (top + dp[adj[v][0]][i]) % mod;

    }



    dp[v][0] = top;



    for(int i = 1; i < adj[v].size(); i++){

        int u = adj[v][i];

        vector<int> pre, preu;



        for(int i = 0; i <= dep[v]; i++){

            pre.push_back(dp[v][i]);

            if(i)

                pre[i] = (pre[i - 1] + pre[i]) % mod;

        }



        for(int j = 0; j <= dep[u]; j++){

            preu.push_back(dp[u][j]);

            if(j)

                preu[j] = (preu[j] + preu[j - 1]) % mod;

        }





        for(int l = dep[v]; l > 0; l--){

            int nerV = min(k - l, l - 1),

                nerU = min(dep[u], min(k - l - 1, l - 1));



            int addend = 1LL * dp[v][l] * preu.back() % mod;

            dp[v][l] = ((nerU >= 0 ? 1LL * dp[v][l] * preu[nerU] % mod : 0) + 

                        (nerV >= 0 && l - 1 <= dep[u] ? 1LL * pre[nerV] * dp[u][l - 1] % mod : 0)) % mod;

            dp[v][l] = (dp[v][l] + addend) % mod;

        }



        dp[v][0] = 1LL * dp[v][0] * preu.back() % mod;

        

    }

}



int main(){

    ios_base::sync_with_stdio(false);

    cin.tie(nullptr);cout.tie(nullptr);



    #ifdef Local

        freopen("C:/Users/Admin/Desktop/Yazilim/C/IO/int.txt","r",stdin);

        freopen("C:/Users/Admin/Desktop/Yazilim/C/IO/out.txt","w",stdout);

    #endif



    cin>>n>>k;



    for(int i = 0; i < n - 1; i++){

        int a, b;

        cin>>a>>b;

        adj[a].push_back(b);

        adj[b].push_back(a);

    }    



    dfs(1, -1);



    /*for(int i = 1; i <= n; i++){

        for(int j = 0; j <= k; j++){

            cout<<dp[i][j]<<" ";

        }

        cout<<endl;

    }*/





    int sum = 0;



    for(int i = 0; i <= k; i++){

        sum = (sum + dp[1][i]) % mod;

    }



    cout<<sum<<endl;



    #ifdef Local

        cout<<endl<<fixed<<setprecision(2)<<1000.0 * clock() / CLOCKS_PER_SEC<< " milliseconds ";

    #endif

}